#include <ntddk.h>

#include "DevMonManager.h"
#include "AutoLock.h"
#include "KDevMon.h"

VOID
DevMonManager::Init(PDRIVER_OBJECT DriverObject)
{
	m_Lock.Init();
	m_DriverObject = DriverObject;
}

NTSTATUS
DevMonManager::AddDevice(PCWSTR Name)
{
	AutoLock<FastMutex> Locker(m_Lock);

	if (m_MonitoredDeviceCount == g_MaxMonitoredDevices) {
		return STATUS_TOO_MANY_NAMES;
	}

	if (FindDevice(Name) >= 0) {
		return STATUS_SUCCESS;
	}

	for (int i = 0; i < g_MaxMonitoredDevices; i++) {
		if (m_Devives[i].m_DeviceObject == nullptr) {
			UNICODE_STRING TargetName;
			RtlInitUnicodeString(&TargetName, Name);

			PFILE_OBJECT   FileObject          = nullptr;
			PDEVICE_OBJECT LowerDeviceObject   = nullptr;
			NTSTATUS Status = IoGetDeviceObjectPointer(&TargetName, FILE_READ_DATA, &FileObject, &LowerDeviceObject);
			if (NT_SUCCESS(Status) == FALSE) {
				KdPrint(("Failed to get device object pointer (%ws) (0x%08X)\n", Name, Status));
				return Status;
			}

			PDEVICE_OBJECT DeviceObject = nullptr;
			WCHAR* Buffer               = nullptr;

			do {
				Status = IoCreateDevice(m_DriverObject, sizeof(DeviceExtension), nullptr, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);
				if (NT_SUCCESS(Status) == FALSE) {
					break;
				}

				Buffer = (WCHAR*)ExAllocatePoolWithTag(PagedPool, TargetName.Length, DRIVER_TAG);
				if (Buffer == nullptr) {
					Status = STATUS_INSUFFICIENT_RESOURCES;
					break;
				}

				DeviceExtension* Ext = (DeviceExtension*)DeviceObject->DeviceExtension;

				DeviceObject->Flags |= LowerDeviceObject->Flags & (DO_BUFFERED_IO | DO_DIRECT_IO);

				DeviceObject->DeviceType                = LowerDeviceObject->DeviceType;
				m_Devives[i].m_DeviceName.Buffer        = Buffer;
				m_Devives[i].m_DeviceName.MaximumLength = TargetName.Length;
				RtlCopyUnicodeString(&m_Devives[i].m_DeviceName, &TargetName);
				m_Devives[i].m_DeviceObject             = DeviceObject;

				Status = IoAttachDeviceToDeviceStackSafe(DeviceObject, LowerDeviceObject, &Ext->m_LowerDeviceObject);
				if (NT_SUCCESS(Status) == FALSE) {
					break;
				}

				m_Devives[i].m_LowerDeviceObject = Ext->m_LowerDeviceObject;
				DeviceObject->Flags &= ~DO_DEVICE_INITIALIZING;
				DeviceObject->Flags |= DO_POWER_PAGABLE;

				m_MonitoredDeviceCount++;
			} while (FALSE);

			if (NT_SUCCESS(Status) == FALSE) {
				if (Buffer != nullptr) {
					ExFreePool(Buffer);
				}
				if (DeviceObject != nullptr) {
					IoDeleteDevice(DeviceObject);
				}
				m_Devives[i].m_DeviceObject = nullptr;
			}

			if (LowerDeviceObject != nullptr) {
				ObDereferenceObject(FileObject);
			}

			return Status;
		}
	}

	return STATUS_UNSUCCESSFUL;
}

BOOLEAN
DevMonManager::RemoveDevice(PCWSTR Name)
{
	AutoLock<FastMutex> Locker(m_Lock);
	int Index = FindDevice(Name);
	if (Index < 0) {
		return FALSE;
	}

	return RemoveDevice(Index);
}

VOID
DevMonManager::RemoveAllDevices()
{
	AutoLock<FastMutex> Locker(m_Lock);
	for (int i = 0; i < g_MaxMonitoredDevices; i++) {
		RemoveDevice(i);
	}
}

int
DevMonManager::FindDevice(PCWSTR Name) CONST
{
	UNICODE_STRING Uname;
	RtlInitUnicodeString(&Uname, Name);

	for (int i = 0; i < g_MaxMonitoredDevices; i++) {
		CONST MonitoredDevice& Device = m_Devives[i];
		if (Device.m_DeviceObject && RtlEqualUnicodeString(&Device.m_DeviceName, &Uname, TRUE) == TRUE) {
			return i;
		}
	}
	return -1;
}

BOOLEAN
DevMonManager::RemoveDevice(int Index)
{
	MonitoredDevice& Device = m_Devives[Index];
	if (Device.m_DeviceObject == nullptr) {
		return FALSE;
	}

	ExFreePool(Device.m_DeviceName.Buffer);
	IoDetachDevice(Device.m_LowerDeviceObject);
	IoDeleteDevice(Device.m_DeviceObject);
	Device.m_DeviceObject = nullptr;
	m_MonitoredDeviceCount--;
	return TRUE;
}